package com.alison.datastream.exercise;

import cn.hutool.core.util.NumberUtil;
import com.alison.datastream.domain.stock.StockPrice;
import com.alison.datastream.domain.stock.StockSource;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.WindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

/**
 * @Author alison
 * @Date 2024/4/6 16:15
 * @Version 1.0
 * @Description
 */
public class E4_1_VWAPJob {

    // 以5分钟为一个时间单位，计算VWAP
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
//        env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);

        env.setParallelism(1);
        DataStream<StockPrice> input = env.addSource(new StockSource("stock/stock-tick-20200108.csv"))
                .assignTimestampsAndWatermarks(WatermarkStrategy.<StockPrice>forMonotonousTimestamps()
                        .withTimestampAssigner(((element, recordTimestamp) -> element.ts)));
        SingleOutputStreamOperator<Tuple2<String, Double>> stream = input.keyBy(x -> x.symbol)
                .window(TumblingEventTimeWindows.of(Time.milliseconds(5)))
                .apply(new WindowFunction<StockPrice, Tuple2<String, Double>, String, TimeWindow>() {
                    @Override
                    public void apply(String key, TimeWindow window, Iterable<StockPrice> input, Collector<Tuple2<String, Double>> out) throws Exception {
                        Double priceSum = 0D, volumnSum = 0D;
                        for (StockPrice stockPrice : input) {
                            priceSum = NumberUtil.add(priceSum.doubleValue(), NumberUtil.mul(stockPrice.price, stockPrice.volume));
                            volumnSum = NumberUtil.add(volumnSum.doubleValue(), stockPrice.volume);
                        }
                        out.collect(Tuple2.of(key, priceSum / volumnSum * 1.0D));
                    }
                });
        stream.print();
        env.execute();
    }
}
